from sklearn.metrics import classification_report, roc_auc_score
from sklearn.model_selection import train_test_split
from shared_functions import *
import xgboost as xgb
import matplotlib.pyplot as plt
import pandas as pd
import time

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = True  # 用来正常显示负号

def train_xgboost(X, y):
    """训练XGBoost模型"""
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    
    print(f"训练集形状: {X_train.shape}, 测试集形状: {X_test.shape}")
    
    # 创建XGBoost分类器
    model = xgb.XGBClassifier(
        n_estimators=100,
        max_depth=6,
        learning_rate=0.3,
        subsample=0.8,
        colsample_bytree=0.8,
        random_state=42,
        # 根据训练集的正例数和反例数计算scale_pos_weight参数
        scale_pos_weight=len(y_train[y_train==0])/len(y_train[y_train==1]) if sum(y_train) > 0 else 1,
        eval_metric='auc'
    )
    
    # 训练模型
    model.fit(
        X_train, y_train,
        eval_set=[(X_test, y_test)],
        # early_stopping_rounds=10,
        verbose=False
    )
    
    return model, X_test, y_test

def evaluate_model(model, X_test, y_test,input_features):
    """评估模型性能"""
    # 预测
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1]
    
    # 计算评估指标
    print("\n分类报告:")
    print(classification_report(y_test, y_pred))
    
    auc_score = roc_auc_score(y_test, y_pred_proba)
    print(f"AUC得分: {auc_score:.4f}")
    
    # 特征重要性
    feature_importance = pd.DataFrame({
        'feature': input_features,
        'importance': model.feature_importances_
    }).sort_values('importance', ascending=False)
    
    print("\n特征重要性:")
    print(feature_importance)
    
    
    # 绘制特征重要性图
    plt.figure(figsize=(10, 6))
    plt.barh(feature_importance['feature'][:10], feature_importance['importance'][:10])
    plt.xlabel('重要性')
    plt.title('Top 10 特征重要性')
    plt.tight_layout()
    plt.savefig('feature_importance.png')
    plt.show()
    
    return auc_score, feature_importance


def main():
    start_time = time.time()
    print("### 开始计时")
    DIR_INPUT = 'simulated-data-transformed/' 

    BEGIN_DATE = "2025-06-11"
    END_DATE = "2025-09-14"

    print("### Load  files")
    transactions_df = read_from_files(DIR_INPUT, BEGIN_DATE, END_DATE)
    print("### {0} transactions loaded, containing {1} fraudulent transactions".format(len(transactions_df),transactions_df.TX_FRAUD.sum()))
    # 输出特征
    output_feature = "TX_FRAUD"
    # 输入特征
    input_features = ['TX_AMOUNT','TX_DURING_WEEKEND', 'TX_DURING_NIGHT', 'CUSTOMER_ID_NB_TX_1DAY_WINDOW',
           'CUSTOMER_ID_AVG_AMOUNT_1DAY_WINDOW', 'CUSTOMER_ID_NB_TX_7DAY_WINDOW',
           'CUSTOMER_ID_AVG_AMOUNT_7DAY_WINDOW', 'CUSTOMER_ID_NB_TX_30DAY_WINDOW',
           'CUSTOMER_ID_AVG_AMOUNT_30DAY_WINDOW', 'TERMINAL_ID_NB_TX_1DAY_WINDOW',
           'TERMINAL_ID_RISK_1DAY_WINDOW', 'TERMINAL_ID_NB_TX_7DAY_WINDOW',
           'TERMINAL_ID_RISK_7DAY_WINDOW', 'TERMINAL_ID_NB_TX_30DAY_WINDOW',
           'TERMINAL_ID_RISK_30DAY_WINDOW']

    print("### Training XGBoost model")
    model, X_test, y_test = train_xgboost(transactions_df[input_features], transactions_df[output_feature])
    auc_score, feature_importance = evaluate_model(model, X_test, y_test,input_features)

    print("### AUC得分:", auc_score)
    model.save_model('xgboost_fraud_detection.model')
    print("### 模型已保存为 'xgboost_fraud_detection.model'")
        
    # 保存特征重要性
    feature_importance.to_csv('feature_importance.csv', index=False)
    print("### 特征重要性已保存为 'feature_importance.csv'")
    print("### 运行时间：{:.2f}秒".format(time.time() - start_time))

